import ext4.data as data


class BlockTail(data.HexArea):
    rules = (
        (data.u32, 1, "checksum"),
    )


class BlockTag3(data.HexArea):
    rules = (
        (data.u32, 1, "blocknr"),
        (data.u32, 1, "flags"),
        (data.u32, 1, "blocknr_high"),
        (data.u32, 1, "checksum"),
    )


class BlockTag(data.HexArea):
    rules = (
        (data.u32, 1, "blocknr"),
        (data.u16, 1, "flags"),
        (data.u16, 1, "blocknr_high"),
        (data.u32, 1, "checksum"),
    )


class BlockType(data.u32):
    codes = (
        "None",
        "DESCRIPTOR",
        "COMMIT",
        "JBD2_SUPERBLOCK_V1",
        "JBD2_SUPERBLOCK_V2",
        "REVOKE",
    )


class Incompat(data.u32):
    bytes = (
        "REVOKE",
        "64BIT",
        "ASYNC_DATA",
        "CSUM_V2",
        "CSUM_V3",
        "FAST_COMMIT",
    )


class Compat(data.u32):
    bytes = ("CHECKSUM",)



class JournalBlockTagFlags(data.u16):
    bytes = (
        "ESCAPE",
        "SAME_UUID",
        "DELETED",
        "LAST_TAG",
    )


class JournalBlockTag3Flags(data.u32):
    bytes = (
        "ESCAPE",
        "SAME_UUID",
        "DELETED",
        "LAST_TAG",
    )


class JournalBlockTagBase(data.HexArea):
    rules = (
        (data.u32, 1, "blocknr"),
        (data.u16, 1, "checksum"),
        (JournalBlockTagFlags, 1, "flags"),
        (data.u32, 1,"blocknr_high"),
    )


class JournalBlockTag3(data.HexArea):
    rules = (
        (data.u32, 1, "blocknr"),
        (data.u32, 1, "flags"),
        (data.u32, 1, "blocknr_high"),
        (data.u32, 1, "checksum"),
    )



class JournalBlockTag(data.Union):
    rules = (
        (JournalBlockTagBase, "tag"),
        (JournalBlockTag3, "tag3"),
    )


class JournalCommitHeader(data.HexArea):
    rules = (
        (data.u32, 1, "magic"),
        (data.u32, 1, "block_type"),
        (data.u32, 1, "sequence"),
        (data.u8, 1, "chksum_type"),
        (data.u8, 1, "chksum_size"),
        (data.u8, 2, "padding"),
        (data.u32, 8, "chksum"),
        (data.Timestamp, 1, "commit_sec"),
        (data.Timestamp, 1, "commit_nsec"),
    )


class JournalHeader(data.HexArea):
    rules = (
        (data.u32, 1, "magic"),
        (BlockType, 1, "blocktype"),
        (data.u32, 1, "sequence"),
    )

    extern_str = ("tags", "uuid")

    def __init__(self, bytes, tagsize=None) -> None:
        super().__init__(bytes)

        # assert self.magic == 0xC03B3998  # ref: JBD2_MAGIC_NUMBER

        if not self.blocktype.has('DESCRIPTOR'):
            return
        assert tagsize is not None

        tags = []
        offset = self.length
        while True:
            tag = JournalBlockTag(self.bytes[offset: offset + tagsize])
            tags.append(tag)

            if (offset == self.length):
                offset += tagsize
                self.uuid = data.IntArray(data.u8, 16, self.bytes[offset: offset + 16])
                """
                ref: jbd2_journal_commit_transaction

                if (first_tag) {                            // after first tag after descriptor tag
                    memcpy (tagp, journal->j_uuid, 16);
                    tagp += 16;
                """
                offset += 16
            else:
                offset += tagsize

            if tag.tag.flags.has("LAST_TAG"):
                break
        
        self.tags = data.Array(tags)


class JournalSurperBlock(data.HexArea):
    rules = (
        (JournalHeader, 1, "header"),
        (data.u32, 1, "blocksize"),
        (data.u32, 1, "maxlen"),
        (data.u32, 1, "first"),
        (data.u32, 1, "sequence"),
        (data.u32, 1, "start"),
        (data.u32, 1, "errno"),
        (Compat, 1, "feature_compat"),
        (Incompat, 1, "feature_incompat"),
        (data.u32, 1, "feature_ro_compat"),
        (data.u8, 16, "uuid"),
        (data.u32, 1, "nr_users"),
        (data.u32, 1, "dynsuper"),
        (data.u32, 1, "max_transaction"),
        (data.u32, 1, "max_trans_data"),
        (data.u8, 1, "checksum_type"),
        (data.u8, 3, "padding2"),
        (data.u32, 1, "num_fc_blks"),
        (data.u32, 41, "padding"),
        (data.u32, 1, "checksum"),
        (data.u8, 16 * 48, "users"),
    )


class Journal:
    def __init__(self, inode) -> None:
        self.inode = inode
        self.jsb = JournalSurperBlock(self.inode.read(JournalSurperBlock.length))

        self.has_csum_v2or3 = self.jsb.feature_incompat.has("CSUM_V2") \
                              or self.jsb.feature_incompat.has("CSUM_V3")

        self.tag_size = self._get_tag_size()

        self.total_len = self.inode.size.value // self.jsb.blocksize
        if self.total_len > self.jsb.maxlen:
            self.total_len = self.jsb.maxlen.value

    def _get_tag_size(self):
        """ref: journal_tag_bytes"""
        if self.jsb.feature_incompat.has("CSUM_V3"):
            return BlockTag3.length

        size = BlockTag.length

        if self.jsb.feature_incompat.has("CSUM_V2"):
            size += 2
        
        if self.jsb.feature_incompat.has("64BIT"):
            return size
        else:
            return size - 4

    def get_journals(self, start=None, cnt=None):
        start = start if start is not None else self.jsb.start.value
        cnt = cnt if cnt is not None else self.total_len
    
        b_nr = start
        while b_nr < start + cnt:
            if b_nr > self.jsb.maxlen:
                b_nr -= self.jsb.maxlen + self.jsb.first

            self.inode.seek(b_nr * self.jsb.blocksize)

            jh = JournalHeader(self.inode.read(self.jsb.blocksize), self._get_tag_size())
            yield jh

            data_block_per_j_block = self.jsb.blocksize // self.inode.ext4.block_size

            if jh.blocktype.has('DESCRIPTOR'):
                """+1 is descriptor block"""
                b_nr += ((len(jh.tags) + 1) // data_block_per_j_block)
            else:
                b_nr += 1
    # def get_journal(self, nr):
    #     pass
